# -*- coding: utf-8 -*-
# pylint: disable=line-too-long
# flake8: noqa: E501
"""The unittests for OpenAI token counter."""
import os
from unittest.async_case import IsolatedAsyncioTestCase

from agentscope.token import OpenAITokenCounter


class OpenAITokenCounterTest(IsolatedAsyncioTestCase):
    """The unittests for the OpenAI token counter."""

    async def asyncSetUp(self) -> None:
        """Set up the test case."""
        self.image_path = os.path.abspath("./image.png")
        self.messages = [
            {
                "role": "system",
                "name": "system",
                "content": [
                    {
                        "type": "text",
                        "text": "You're a helpful assistant.",
                    },
                ],
            },
            {
                "role": "user",
                "name": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "What is the capital of France?",
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": "",
                        },
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": "https://help-static-aliyun-doc.aliyuncs.com/file-manage-files/zh-CN/20241022/emyrja/dog_and_girl.jpeg",
                        },
                    },
                ],
            },
            {
                "role": "assistant",
                "name": "assistant",
                "content": [
                    {
                        "type": "text",
                        "text": "The capital of France is Paris.",
                    },
                ],
            },
            {
                "role": "user",
                "name": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "What is the capital of Germany?",
                    },
                ],
            },
            {
                "role": "assistant",
                "name": "assistant",
                "content": [
                    {
                        "type": "text",
                        "text": "The capital of Germany is Berlin.",
                    },
                ],
            },
            {
                "role": "user",
                "name": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "What is the capital of Japan?",
                    },
                ],
            },
            {
                "role": "assistant",
                "name": "assistant",
                "content": None,
                "tool_calls": [
                    {
                        "id": "1",
                        "type": "function",
                        "function": {
                            "name": "get_capital",
                            "arguments": '{"country": "Japan"}',
                        },
                    },
                ],
            },
            {
                "role": "tool",
                "tool_call_id": "1",
                "content": "The capital of Japan is Tokyo.",
            },
            {
                "role": "assistant",
                "name": "assistant",
                "content": [
                    {
                        "type": "text",
                        "text": "The capital of Japan is Tokyo.",
                    },
                ],
            },
        ]
        self.tools = [
            {
                "type": "function",
                "function": {
                    "name": "get_location",
                    "description": "Get the location of the user",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "city": {
                                "type": "string",
                                "description": "The city to get the location "
                                "for",
                            },
                        },
                        "required": ["city"],
                    },
                },
            },
        ]

    async def test_openai_token_counter(self) -> None:
        """Test the OpenAI token counter."""
        counter = OpenAITokenCounter(
            model_name="gpt-4o-mini",
        )
        n_tokens = await counter.count(self.messages)
        self.assertEqual(n_tokens, 2016)

        n_tokens = await counter.count(self.messages, self.tools)
        self.assertEqual(n_tokens, 2058)

        counter = OpenAITokenCounter(
            model_name="o3-mini",
        )
        n_tokens = await counter.count(self.messages)
        self.assertEqual(n_tokens, 1796)

        n_tokens = await counter.count(self.messages, self.tools)
        self.assertEqual(n_tokens, 1841)
